import argparse
import pysam
import re
from collections import Counter
import csv
from deepsignal3.utils.process_utils import complement_seq

def parse_args():
    parser = argparse.ArgumentParser(description="Count motifs in BAM file and find motif reference positions not in TSV with 6mer sequences")
    parser.add_argument("-b", "--bam", required=True, help="Input BAM file")
    parser.add_argument("-m", "--motif", required=True, help="Motif pattern to search for")
    parser.add_argument("-t", "--tsv", required=True, help="TSV file from modkit extract calls")
    return parser.parse_args()

def load_tsv_ref_positions(tsv_file):
    """Load all reference positions from TSV file"""
    tsv_ref_positions = set()
    with open(tsv_file, 'r') as f:
        reader = csv.DictReader(f, delimiter='\t')
        for row in reader:
            read_id = row['read_id']
            ref_pos = int(row['ref_position'])
            if ref_pos >= 0:  # Only include mapped positions
                tsv_ref_positions.add((read_id, ref_pos))
    return tsv_ref_positions

def count_motif_in_sequence(sequence, motif):
    """Count occurrences of motif in a sequence using regex"""
    pattern = re.compile(motif, re.IGNORECASE)
    return len([match.start() for match in pattern.finditer(sequence)])

def get_6mer(sequence, start_pos, seq_type="read"):
    """Extract 6mer centered at start_pos (0-based), pad with N if out of bounds"""
    left = start_pos - 3
    right = start_pos + 3
    warning = None
    if left < 0:
        pad_left = -left
        sequence = "N" * pad_left + sequence
        left += pad_left
        warning = f"6mer truncated at {seq_type} start, padded with N"
    if right > len(sequence):
        pad_right = right - len(sequence)
        sequence = sequence + "N" * pad_right
        warning = f"6mer truncated at {seq_type} end, padded with N" if warning is None else warning + ", truncated at end"
    return sequence[left:right], warning

def process_bam(bam_file, motif, tsv_ref_positions):
    """Process BAM file, count motifs, and find motif ref positions not in TSV with 6mers"""
    try:
        bam = pysam.AlignmentFile(bam_file, "rb")
    except ValueError as e:
        print(f"Error opening BAM file: {e}")
        return None, None, None, None

    orig_counts = Counter()
    aligned_counts = Counter()
    motif_locations = []
    warnings = []
    bam_ref_positions = set()

    for read in bam:
        if read.is_unmapped:
            continue
            
        # Original sequence motif counts
        orig_seq = read.query_sequence.upper()  # Assume get_forward_sequence is query_sequence
        orig_counts[count_motif_in_sequence(orig_seq, motif)] += 1

        # Aligned sequence motif counts
        aligned_seq = read.get_reference_sequence().upper() if read.has_tag("MD") else orig_seq
        aligned_counts[count_motif_in_sequence(aligned_seq, motif)] += 1

        # Find motifs and check reference positions
        read_id = read.query_name
        ref_loc = read.get_reference_positions(full_length=True)
        seq_len = len(orig_seq)
        non_none_ref_loc = [p for p in ref_loc if p is not None]
        if not non_none_ref_loc:
            continue
        R_start = min(non_none_ref_loc)
        R_end = max(non_none_ref_loc)
        for match in re.finditer(motif, orig_seq, re.IGNORECASE):
            start = match.start()
            end = match.end()
            # Ensure motif is fully mapped
            if all(ref_pos is not None for ref_pos in ref_loc[start:end]):
                motif_ref_positions = [ref_pos for ref_pos in ref_loc[start:end]]
                # Add to bam_ref_positions for counting
                for ref_pos in motif_ref_positions:
                    bam_ref_positions.add((read_id, ref_pos))
                # Check if any reference position is not in TSV
                if any((read_id, ref_pos) not in tsv_ref_positions for ref_pos in motif_ref_positions):
                    chrom = read.reference_name
                    # Map read start to reference position with strand correction
                    P = ref_loc[start] if read.is_forward else ref_loc[seq_len - start - 1]
                    if P is None:
                        continue
                    i = P - R_start
                    # Validate index
                    if not (0 <= i < len(aligned_seq)):
                        warnings.append(f"{read_id} at read_pos {start}: Invalid reference index for 6mer")
                        continue
                    # Get read 6mer
                    read_6mer, read_warning = get_6mer(orig_seq, start, seq_type="read")
                    if read_warning:
                        warnings.append(f"{read_id} at read_pos {start}: {read_warning}")
                    # Get reference 6mer
                    ref_6mer, ref_warning = get_6mer(aligned_seq, i, seq_type="ref")
                    if ref_warning:
                        warnings.append(f"{read_id} at read_pos {start}: {ref_warning}")
                    # Record motif location
                    motif_locations.append((
                        read_id, start, end-1, chrom, P, max(motif_ref_positions), read_6mer, ref_6mer
                    ))

    bam.close()
    return orig_counts, aligned_counts, motif_locations, warnings, bam_ref_positions

def print_output(orig_counts, aligned_counts, motif_locations, motif, warnings, tsv_ref_positions, bam_ref_positions):
    """Print motif counts, motif locations, 6mer sequences, and debug info"""
    print(f"Motif: {motif}\n")
    
    print("Debug Info:")
    print(f"TSV reference positions: {len(tsv_ref_positions)}")
    print(f"BAM motif reference positions: {len(bam_ref_positions)}")
    
    print("\nOriginal Sequence Motif Counts:")
    print("Count\tNumber of Reads")
    for count, num_reads in sorted(orig_counts.items()):
        print(f"{count}\t{num_reads}")
        
    print("\nAligned Sequence Motif Counts:")
    print("Count\tNumber of Reads")
    for count, num_reads in sorted(aligned_counts.items()):
        print(f"{count}\t{num_reads}")
        
    print("\nMotif Reference Positions with Any Ref Pos Not in TSV (1-based):")
    if motif_locations:
        print("Read ID\tRead Start\tRead End\tChrom\tRef Start\tRef End\tRead 6mer\tRef 6mer")
        for loc in sorted(motif_locations, key=lambda x: (x[3], x[4])):
            print(f"{loc[0]}\t{loc[1]}\t{loc[2]}\t{loc[3]}\t{loc[4]}\t{loc[5]}\t{loc[6]}\t{loc[7]}")
    else:
        print("No motif reference positions found with any ref pos not in TSV.")
    
    if warnings:
        print("\nWarnings:")
        for warning in warnings:
            print(f"- {warning}")

def main():
    args = parse_args()
    tsv_ref_positions = load_tsv_ref_positions(args.tsv)
    orig_counts, aligned_counts, motif_locations, warnings, bam_ref_positions = process_bam(
        args.bam, args.motif, tsv_ref_positions
    )
    
    if orig_counts is None:
        print("Failed to process BAM file")
        return
        
    print_output(orig_counts, aligned_counts, motif_locations, args.motif, warnings, tsv_ref_positions, bam_ref_positions)

if __name__ == "__main__":
    main()